// Copyright (C) 2010 and onwards Harik Shazeer Labs
// 
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// 
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.
//
// Authors: Noam Shazeer and Georges Harik

#ifndef _ALLOCATOR_H_
#define _ALLOCATOR_H_

#include <stdarg.h>
#include <pthread.h>
#include "util.h"
#include "objects.h"
#include "jobqueue.h"

// When a thread allocator asks the global allocator for memory, 
// it gets an allocation block with a start point, and a limit point. 

// Approximately how many words total we can allocate
extern const uint64 G_MAX_PHYSICAL_MEMORY;

// number of allocation blocks
extern const uint64 G_NUM_ALLOCATION_BLOCKS;

// virtual size of an allocation block (in words) 
extern const uint64 G_ALLOCATION_BLOCK_SIZE;

// when the global allocator gives a block to the thread allocator,
// make sure the limit point minus the start point is >= this (words)
extern const uint64 G_ALLOCATION_SIZE; 

// we need to know the page size when calling mmap in words
extern const uint64 G_PAGE_SIZE;

const uint64 KILL_THREADS_URGENCY = 0;
const uint64 DELETE_TEMPORARY_ADS_URGENCY = 1;
const uint64 KILL_PROCESS_URGENCY = 2;

// Allocators don't know about specific objects
class Frontier;

class AllocatorDependant {
 public:
  virtual int32 NumShards() = 0;
  virtual void FlagReferences(int32 shard, Frontier * frontier) = 0;
  virtual void ObjectsDeleted(int32 shard) = 0;
  virtual void DebugDependant() {};

  AllocatorDependant();
  virtual ~AllocatorDependant();
};

struct ThreadAllocator;

struct AllocationBlock {
  uint64 * data_;
  uint64 * next_;
  // How far this block is allowed to allocate
  uint64 * limit_;
  uint64 * new_next_; // used in garbage collection
  ThreadAllocator * owner_;
  uint64 AllocatedSize() {return next_ - data_;}
  uint64 AvailableSize() { return limit_ - next_;}  
  AllocationBlock();
  ~AllocationBlock();
  // set next_ and limit_ to new_next. 
  // release physical memory back to the system.
  // only called during garbage collection.
  void Shrink();
} __attribute__ ((aligned (64)));

// A set of Objects to be expanded
// The set is a union of 2 types of Object sets
// Object set 1 is a linked link using the link__ field
// Object set 2 is a range inside a big array

struct Frontier {
Frontier() : head_(NULL), list_expansion_size_(0), 
    start_(NULL), end_(NULL), seen_memory_(0) {}
  void AddObject(Object * o);

  // Add the range of Objects into the linked list
  void ExpandRange(Object ** begin, Object ** end);

  void Expand();

  // For the linked list
  Object * head_;
  // Sum of the array sizes in the linked list
  uint64 list_expansion_size_;

  // For the range
  Object ** start_;
  Object ** end_;

  // How much memory this frontier has seen
  uint64 seen_memory_;
};

struct GlobalAllocator : public AllocatorDependant {
  GlobalAllocator();

  void GarbageCollect(uint64 nwords);
  void ComputeReachable();
  bool HasEnoughMemory(uint64 nwords);
  void RewriteArrays();
  void NotifyObjectsDeleted();
  void MoveObjects();

  // Garbage Collection routines
  int32 NumShards() { return 100; }
  void FlagReferences(int32 shard, Frontier * frontier);
  void ObjectsDeleted(int32 shard);
  void AddDependant(AllocatorDependant *d) {
    dependants_.insert(d);
  }
  void AddThreadAllocator(ThreadAllocator *ta) {
    thread_allocators_.push_back(ta);
  }

  // Change where the block is allowed to write until
  void SetBlockLimit(AllocationBlock *b, uint64 * new_limit) {
    memory_allocated_ += new_limit - b->limit_;
    b->limit_ = new_limit;
  }

  // called internally - reshelve a block. 
  void RecallBlock(AllocationBlock *b);

  // give the thread a block with a suitable limit. 
  void Allocate(ThreadAllocator *ta, uint64 needed);

  bool BlockIsUsable(AllocationBlock *b) {
    // TODO: load balance by skipping overfull blocks
    return !b->owner_;
  }
  
  // to avoid creating duplicate arrays
  Object ** array_cache_;
  
  // How much memory have we given the ThreadAllocators permission to use.
  //  ( an upper bound on the physical memory in use )
  uint64 memory_allocated_;

  // Used in garbage collection to tell if we can stop
  uint64 memory_used_;

  // These are involved in flagging memory usage
  set<AllocatorDependant *> dependants_;

  // A list of ThreadAllocators
  vector<ThreadAllocator *> thread_allocators_;
  
  // These are all allocated objects globally
  // Consistent with thread allocators at the start & end of GC
  // vector<Object*> allocated_;
  AllocationBlock * blocks_;

  // we round robin giving out blocks
  uint32 next_block_;

  // To tell if following importance is done
  bool follow_failure_;
  
  // To synchronize garbage collection
  pthread_mutex_t lock_;

  // In garbage collection, a job for propagating reachability
  JobQueue::Job *propagate_job_;

  // This value increases as we try to garbage collect more stuff
  uint64 urgency_;
};

extern GlobalAllocator G_ALLOC;

struct ThreadAllocatorException {
};

// One of these exists per thread that runs a virtual machine
struct ThreadAllocator : public AllocatorDependant {
  ThreadAllocator();
  
  // allocate raw memory - length is in words
  uint64 * Allocate(uint64 length);

  Object* AllocateArray(uint64 length);
  // Object* AllocateByteArray(uint64 length);

  Object * MakeArray(Object ** contents, uint64 length);

  Object * MakeList(Object ** contents, uint64 length);

  // Make an immutable array of what's passed in after count
  // Must be count arguments passed in
  Object * MakeArray(int32 count, ...) {
    CHECK(count < 100);
    va_list va;
    va_start(va, count);
    for (int32 i=0; i<count; i++)
      temp_[i] = *va_arg(va, Object **);
    return MakeArray(temp_, count);
  }
  
  // Make an immutable array based on a vector of Object *s
  Object * MakeArray(vector<Object *>& v) {
    if (v.size() == 0) return MakeArray((Object **)NULL, 0);
    return MakeArray(&(v[0]), v.size());
  }

  // Make an immutable array based on a vector of Object *s
  Object * MakeArray(StableVector<Object *>& v) {
    return MakeArray(&(v[0]), v.size());
  }

  // Flagging references for garbage collection
  void FlagReferences(int32 shard, Frontier * frontier);
  void ObjectsDeleted(int32 shard);
  int32 NumShards() { return 1; } 

  // Acquiring allocator locks / releasing them
  void AcquireLockCheckException() {
    pthread_mutex_lock(&lock_);
    CheckException();
  }

  void AcquireLock() {
    pthread_mutex_lock(&lock_);
  }
  
  void CheckException() {
    if (throw_exception_) {
      throw_exception_ = false;
      throw (ThreadAllocatorException());
    }
  }

  void ReleaseLock() {
    pthread_mutex_unlock(&lock_);
  }

  // These are temporarily referenced externally
  // TODO: may want to make this a linked list of ObjectSavers. 
  StableVector<Object **> saved_;
  vector<vector<Object *> *> saved_vectors_;

  // These are all allocated objects
  AllocationBlock * block_; 

  bool throw_exception_;

  // temporary space for making small arrays
  Object * temp_[100];

  // statistics:
  uint64 array_cache_tries_;
  uint64 array_cache_hits_;

  // Lock must be held while the thread owning this ThreadAllocator runs
  // This allows garbage collection to run
  // We release this lock before taking any potentially blocking operation
  // We reacquire after a blocking operation
 private:
  pthread_mutex_t lock_;
};

// Save an object from garbage collection
struct ObjectSaver {
  ObjectSaver(ThreadAllocator *a, Object ** o) 
  { a_ = a; a_->saved_.push_back(o); }
  ~ObjectSaver() 
  { a_->saved_.pop_back(); }

  ThreadAllocator *a_;
};

// Save a vector of objects from garbage collection
struct ObjectVectorSaver {
  ObjectVectorSaver(ThreadAllocator *a, vector<Object *> * v) {
    a_ = a;
    a_->saved_vectors_.push_back(v); }
  ~ObjectVectorSaver() { a_->saved_vectors_.pop_back(); }
  ThreadAllocator *a_;
};

extern string ToString(Object*);

// Run an object handler on all objects in a block
struct BlockTask : public JobQueue::Task {
  virtual void Handle(Object * o) {};
  virtual void GlobalStart() {}
  virtual void StartBlock() {}
  virtual void FinishBlock() {}

  void Run() {
    StartBlock();
    Object * o = (Object *) b_->data_;
    for (; o < (Object *) b_->next_;) {
      Object * next = NextObject(o);
      if (GetType(o) != FREE_SPACE) Handle(o);
      o = next;
    }
    FinishBlock();
  }

  // Computed internally
  AllocationBlock * b_;
};

// BT is a subclass of BlockTask
template <class BT> 
void RipAllObjects() {
  JobQueue::Job j(G_JQ);
  for (uint32 c=0; c<G_NUM_ALLOCATION_BLOCKS; c++) {
    BT * bt = new BT();
    if (c == 0) bt->GlobalStart();
    bt->b_ = G_ALLOC.blocks_ + c;
    j.AddTask(bt);
  }
  j.Wait();
}

// This allows an array to be used as a vector
// TODO: The array had better be allocated mutable
// The first element stores its length
// The remaining elements actually store the vector
// A null pointer can represent an empty vector

void VectorObjectClear(Object *& v);
uint64 VectorObjectSize(Object * v);
void VectorObjectSetSize(Object * v, uint64 size);
Object *& VectorObjectRef(Object *v, uint64 pos);

// pushes increment NULLs onto the end of the VectorObject.  
// If the VectorObject is garbage collected, returns false.
// This is thread unsafe. You need to make sure only 1 thread runs this.
bool VectorObjectExtend(ThreadAllocator *a, 
			Object *& v, 
			uint64 increment);

// If the vector you are pointing to gets garbage collected during
// resizing to add an element, it just adds the element to an empty vector
// This is thread unsafe. You need to make sure only 1 thread runs this.
void VectorObjectPushBack(ThreadAllocator *a, 
			  Object *& v,
			  Object ** val,
			  uint64 length,
			  uint64 num_copies);

void VectorObjectParallelPush
(ThreadAllocator *a,
 Object *& v,
 const vector<pair<pair<uint64, Object *> *, 
 pair<uint64, Object *> *> > & ranges);
			      
#endif
